/* 
* Copyright (C) 2010, John Asmuth

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
* 
*  $Revision$
*  $Date$
*  $Author$
*  $HeadURL$
* 
*/

package rlglue

import (
	"bytes"
	"errors"
	"fmt"
	"math"
	"strconv"
	"strings"
)

type IntRange struct {
	Min, Max int32
}

func (ir IntRange) Span() int32 {
	return ir.Max - ir.Min
}

func (ir IntRange) Count() (count uint64) {
	count = uint64(ir.Max)
	if ir.Min < 0 {
		count += uint64(-1 * ir.Min)
	} else {
		count -= uint64(ir.Min)
	}
	count += 1
	return
}

func (ir IntRange) Index(val int32) uint64 {
	return uint64(val - int32(ir.Min))
}

func (ir IntRange) Value(index uint64) int32 {
	return int32(index + uint64(ir.Min))
}

type DoubleRange struct {
	Min, Max float64
}

func (dr DoubleRange) Span() float64 {
	return dr.Max - dr.Min
}

type IntRanges []IntRange
type DoubleRanges []DoubleRange

type TypeRanges struct {
	Ints    IntRanges
	Doubles DoubleRanges
	Chars   int32
}

func (rs IntRanges) Count() (count uint64) {
	count = 1
	for _, ir := range rs {
		count *= ir.Count()
	}
	return
}

func (rs IntRanges) Index(values []int32) (index uint64) {
	var base uint64 = 1
	for i, val := range values {
		index += rs[i].Index(val) * base
		base *= rs[i].Count()
	}
	return
}

func (rs IntRanges) Values(index uint64) (values []int32) {
	values = make([]int32, len(rs))
	for i, _ := range rs {
		base := rs[i].Count()
		values[i] = rs[i].Value(index % base)
		index = index / base
	}
	return
}

type TaskSpec struct {
	ProblemType    string
	DiscountFactor float64
	Obs            TypeRanges
	Act            TypeRanges
	Reward         DoubleRange
	Extra          string
}

func (this *TaskSpec) String() (res string) {
	buf := bytes.NewBuffer([]byte{})

	var err error

	//"VERSION RL-Glue-3.0 PROBLEMTYPE episodic DISCOUNTFACTOR %f OBSERVATIONS INTS (%d 0 1) ACTIONS INTS (0 %d) REWARDS (-1.0 1.0)"
	wr := func(format string, args ...interface{}) {
		if err == nil {
			_, err = fmt.Fprintf(buf, format, args...)
		}
	}

	rangesFoo := func(tr TypeRanges) {
		if len(tr.Ints) > 0 {
			wr("INTS ")
			for _, ir := range tr.Ints {
				wr("(%d %d) ", ir.Min, ir.Max)
			}
		}
		if len(tr.Ints) > 0 {
			wr("DOUBLES ")
			for _, dr := range tr.Doubles {
				wr("(%f %f) ", dr.Min, dr.Max)
			}
		}
		if tr.Chars != 0 {
			wr("CHARCOUNT %d ", tr.Chars)
		}
	}

	wr("VERSION RL-Glue-3.0 ")

	wr("PROBLEMTYPE %s ", this.ProblemType)

	wr("DISCOUNTFACTOR %f ", this.DiscountFactor)

	wr("OBSERVATIONS ")
	rangesFoo(this.Obs)
	wr("ACTIONS ")
	rangesFoo(this.Act)

	wr("REWARDS (%f %f) ", this.Reward.Min, this.Reward.Max)

	if this.Extra != "" {
		wr("EXTRA %s ", this.Extra)
	}

	if err != nil {
		return err.Error()
	}

	return buf.String()
}

type tokenlist struct {
	s string
}

func (tl *tokenlist) Empty() bool {
	return tl.s == ""
}

func (tl *tokenlist) Peek() (res string) {
	if tl.s[0] == '(' {
		i := strings.Index(tl.s, ")")
		if i == -1 {
			i = len(tl.s)
		} else {
			res = tl.s[0 : i+1]
		}
		return
	}
	i := strings.Index(tl.s, " ")
	if i == -1 {
		res = tl.s
	} else {
		res = tl.s[0:i]
	}
	return
}

func (tl *tokenlist) Next() (res string) {
	res = tl.Peek()
	//fmt.Printf("Token:%s\n", res)
	tl.s = strings.TrimLeft(tl.s[len(res):len(tl.s)], " \t\n")
	return
}

func (tl *tokenlist) Rest() (res string) {
	res = tl.s
	tl.s = ""
	return
}

func atoi32(s string) (res int32, err error) {
	val, err := strconv.Atoi(s)
	res = int32(val)
	return
}

func atof64(s string) (res float64, err error) {
	if s == "NEGINF" {
		res = math.Inf(-1)
		return
	}
	if s == "POSINF" {
		res = math.Inf(1)
		return
	}
	res, err = strconv.ParseFloat(s, 64)
	return
}

func parseRanges(tokens *tokenlist) (intRanges []IntRange, doubleRanges []DoubleRange, charCount int32, err error) {
	for {
		switch tokens.Peek() {
		case "INTS":
			tokens.Next()
			intRangeV := make([]interface{}, 0, 5)
			for tokens.Peek()[0] == '(' {
				rangeStr := tokens.Next()
				intStrs := strings.Split(rangeStr[1:len(rangeStr)-1], " ")
				if len(intStrs) == 2 {
					var r IntRange
					r.Min, err = atoi32(intStrs[0])
					r.Max, err = atoi32(intStrs[1])
					intRangeV = append(intRangeV, &r)
				}
				if len(intStrs) == 3 {
					var count int
					count, err = strconv.Atoi(intStrs[0])
					for i := 0; i < count; i++ {
						var r IntRange
						r.Min, err = atoi32(intStrs[1])
						r.Max, err = atoi32(intStrs[2])
						intRangeV = append(intRangeV, &r)
					}
				}
			}
			intRanges = make([]IntRange, len(intRangeV))
			for index, ir := range intRangeV {
				intRanges[index] = *ir.(*IntRange)
			}
		case "DOUBLES":
			tokens.Next()
			doubleRangeV := make([]interface{}, 0, 5)
			for tokens.Peek()[0] == '(' {
				rangeStr := tokens.Next()
				doubleStrs := strings.Split(rangeStr[1:len(rangeStr)-1], " ")
				if len(doubleStrs) == 2 {
					var r DoubleRange
					r.Min, err = atof64(doubleStrs[0])
					r.Max, err = atof64(doubleStrs[1])
					doubleRangeV = append(doubleRangeV, &r)
				}
				if len(doubleStrs) == 3 {
					var count int
					count, err = strconv.Atoi(doubleStrs[0])
					for i := 0; i < count; i++ {
						var r DoubleRange
						r.Min, err = atof64(doubleStrs[1])
						r.Max, err = atof64(doubleStrs[2])
						doubleRangeV = append(doubleRangeV, &r)
					}
				}
			}
			doubleRanges = make([]DoubleRange, len(doubleRangeV))
			for index, ir := range doubleRangeV {
				doubleRanges[index] = *ir.(*DoubleRange)
			}
		case "CHARCOUNT":
			tokens.Next()
			charCount, err = atoi32(tokens.Next())
		default:
			return
		}
	}
	return
}

func ParseTaskSpec(taskString string) (ts *TaskSpec, err error) {
	tokens := &tokenlist{taskString}

	if tokens.Next() != "VERSION" || tokens.Next() != "RL-Glue-3.0" {
		// not what we're ready for
		err = errors.New("Incorrect version")
		return
	}

	ts = new(TaskSpec)
	ts.DiscountFactor = 1

	for !tokens.Empty() {
		switch tokens.Next() {
		case "DISCOUNTFACTOR":
			if tokens.Empty() {
				err = errors.New("No discount factor")
				return
			}
			ts.DiscountFactor, err = atof64(tokens.Next())
			if err != nil {
				return
			}
		case "PROBLEMTYPE":
			ts.ProblemType = tokens.Next()
		case "OBSERVATIONS":
			ts.Obs.Ints, ts.Obs.Doubles, ts.Obs.Chars, err = parseRanges(tokens)
		case "ACTIONS":
			ts.Act.Ints, ts.Act.Doubles, ts.Act.Chars, err = parseRanges(tokens)
		case "REWARDS":
			rstr := tokens.Next()
			rstrs := strings.Split(rstr[1:len(rstr)-1], " ")
			ts.Reward.Min, err = atof64(rstrs[0])
			ts.Reward.Max, err = atof64(rstrs[1])
		case "EXTRA":
			ts.Extra = tokens.Rest()
		}
	}
	return
}
